/**
 * CS255 project 2
 */
package mitm;

import iaik.security.rsa.RSAPrivateKey;

import java.io.*;
import java.net.*;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.Signature;

import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocket;
import javax.net.ssl.TrustManager;

public class MITMAdminClient
{
    private boolean m_useChallengeResponse = false;
    private Socket m_remoteSocket;
    private String password;
    private String command;
    private String commonName = "";
    private PrivateKey m_privateKey;
    
    public static void main( String [] args ) {
	MITMAdminClient admin = new MITMAdminClient( args );
	admin.run();
    }

     private Error printUsage() {
	System.err.println(
	    "\n" +
	    "Usage: " +
	    "\n java " + MITMAdminClient.class + " <options>" +
	    "\n" +
	    "\n Where options can include:" +
	    "\n" +
	    "\n   <-password <pass> >           Used as keystore password if keystore specified" +
	    "\n   <-keystore <keystore> >       Keystore for challenge response authentication" +
	    "\n         Note: Challenge response authentication used if keystore specified" +
	    "\n         Default is password-only authentication" +
	    "\n   <-cmd <shudown|stats>" +
	    "\n   [-remoteHost <host name/ip>]  Default is localhost" +
	    "\n   [-remotePort <port>]          Default is 8002" +
	    "\n"
	    );

	System.exit(1);
	return null;
    }


    private static class TrustEveryone implements javax.net.ssl.X509TrustManager
    {
	public void checkClientTrusted(java.security.cert.X509Certificate[] chain,
				       String authenticationType) {
	}

	public void checkServerTrusted(java.security.cert.X509Certificate[] chain,
				       String authenticationType) {
	}

	public java.security.cert.X509Certificate[] getAcceptedIssuers()
	{
	    return null;
	}
    }


    private MITMAdminClient( String [] args ) {
	int remotePort = 8002;
	String remoteHost = "localhost";
	String keystoreFile = "";
		
	if( args.length < 3 )
	    throw printUsage();
	
	try {
	    for (int i=0; i<args.length; i++)
	    {
		if (args[i].equals("-remoteHost")) {
		    remoteHost = args[++i];
		} else if (args[i].equals("-remotePort")) {
		    remotePort = Integer.parseInt(args[++i]);
		} else if (args[i].equals("-password")) {
		    password = args[++i];
		} else if (args[i].equals("-cmd")) {
		    command = args[++i];
		    if( command.equals("enable") || command.equals("disable") ) {
			commonName = args[++i];
		    }
		} else if (args[i].equals("-keystore")) {
            m_useChallengeResponse = true;
            keystoreFile = args[++i];
        }
        else {
		    throw printUsage();
		}
	    }

	    SSLContext sslContext = SSLContext.getInstance( "SSL" );

	    sslContext.init(
    		new javax.net.ssl.KeyManager[] {}
    		, new TrustManager[] { new TrustEveryone() }
    		, null
    		);
	    
	    if (m_useChallengeResponse)
	    {
		    try {
	            // Open the file that is the first
	            // command line parameter
	
	
	            final KeyStore keyStore;
	
	            keyStore = KeyStore.getInstance("jks");
	            keyStore.load(new FileInputStream(keystoreFile), password.toCharArray());
	
	            // Get our key pair and our own DN (not the remote server's DN) from the keystore.
	
	            m_privateKey = (PrivateKey)keyStore.getKey("mykey", password.toCharArray()) ;
	
	        }catch (Exception e){//Catch exception if any
	            System.err.println("Error: failed to load the keystore...bad password");
	            System.exit(1);
	        }
	    }

	    m_remoteSocket = (SSLSocket) sslContext.getSocketFactory().createSocket( remoteHost, remotePort );
	}
	catch (Exception e) {
	    throw printUsage();
	}

    }
    
    public void run() 
    {
        try
        {
        	if (m_useChallengeResponse)
	        {
	            if( m_remoteSocket != null )
	            {
	                PrintWriter writer = new PrintWriter( m_remoteSocket.getOutputStream() );

	                //format the client command, send blank password if in password mode
	                writer.print("requestChallenge command:" + command);
	                writer.flush();
	
	                BufferedReader r = new BufferedReader(new InputStreamReader(m_remoteSocket.getInputStream()));
	                String line = null;
	
	                while ((line = r.readLine()) != null)
	                {
	                    //check for a challenge, ignore it if we are not in cr mode
	                    String[] challenge = line.split(":");
	
	                    if(challenge.length == 2 && challenge[0].equals("challenge"))
	                    {
	                        //get the challenge string
	                        String random = challenge[1].trim();
	
	                        //create our signature method
	                        Signature s = Signature.getInstance("SHA1withRSA");
	                        s.initSign(m_privateKey);
	
	                        //concatenate our challenge string with our command string to prevent tampering
	                        byte[] randomBytes = javax.xml.bind.DatatypeConverter.parseBase64Binary(random);
	                        byte[] commandBytes =  command.getBytes();
	
	                        byte[] temp = new byte[commandBytes.length + randomBytes.length];
	
	                        System.arraycopy(commandBytes, 0, temp, 0, commandBytes.length);
	                        System.arraycopy(randomBytes, 0, temp,  commandBytes.length, randomBytes.length);
	
	                        //add our data to the signer
	                        s.update(temp);
	
	                        //sign it and return as a base64 string
	                        String response = javax.xml.bind.DatatypeConverter.printBase64Binary(s.sign());
	
	                        System.out.println("Responding to challenge");
	
	                        writer.print("response:" +response );
	                        writer.flush();
	                    }
	                    //otherwise its just a status message
	                    else
	                    {
	                        System.out.println(line);
	                        System.exit(0);
	                    }
	                }
	            }
	        }
        	else
        	{
        		if( m_remoteSocket != null ) {
        			PrintWriter writer =
        			    new PrintWriter( m_remoteSocket.getOutputStream() );
        			writer.println("password:"+password);
        			writer.println("command:"+command);
        			writer.println("CN:"+commonName);
        			writer.flush();
        		}

        		    // now read back any response

    		    System.out.println("");
    		    System.out.println("Receiving input from MITM proxy:");
    		    System.out.println("");
    		    BufferedReader r = new BufferedReader(new InputStreamReader(m_remoteSocket.getInputStream()));
    		    String line = null;
    		    while ((line = r.readLine()) != null) {
    		    	System.out.println(line);
        		}
        	}
        }
        catch (Exception e)
        {
            e.printStackTrace();
        }

        //this does not indicate an error...not sure why starter code sends it to stderr
        System.err.println("Admin Client exited");

        System.exit(0);
    }
}
